{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b518b04cbfe0"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "906e07f6e562"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e2d97c7e31aa"
      },
      "source": [
        "# Making new Layers and Models via subclassing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4e352274064f"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/keras/custom_layers_and_models\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/snapshot-keras/site/en/guide/keras/custom_layers_and_models.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/keras-team/keras-io/blob/master/guides/making_new_layers_and_models_via_subclassing.py\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/keras/custom_layers_and_models.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8d4ac441b1fc"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4e7dce39dd1d"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow import keras"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7b363673d96c"
      },
      "source": [
        "## The `Layer` class: the combination of state (weights) and some computation\n",
        "\n",
        "One of the central abstraction in Keras is the `Layer` class. A layer\n",
        "encapsulates both a state (the layer's \"weights\") and a transformation from\n",
        "inputs to outputs (a \"call\", the layer's forward pass).\n",
        "\n",
        "Here's a densely-connected layer. It has a state: the variables `w` and `b`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "59b8317dbd3c"
      },
      "outputs": [],
      "source": [
        "class Linear(keras.layers.Layer):\n",
        "    def __init__(self, units=32, input_dim=32):\n",
        "        super(Linear, self).__init__()\n",
        "        w_init = tf.random_normal_initializer()\n",
        "        self.w = tf.Variable(\n",
        "            initial_value=w_init(shape=(input_dim, units), dtype=\"float32\"),\n",
        "            trainable=True,\n",
        "        )\n",
        "        b_init = tf.zeros_initializer()\n",
        "        self.b = tf.Variable(\n",
        "            initial_value=b_init(shape=(units,), dtype=\"float32\"), trainable=True\n",
        "        )\n",
        "\n",
        "    def call(self, inputs):\n",
        "        return tf.matmul(inputs, self.w) + self.b\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dac8fb03a642"
      },
      "source": [
        "You would use a layer by calling it on some tensor input(s), much like a Python\n",
        "function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cdcd15d5e68a"
      },
      "outputs": [],
      "source": [
        "x = tf.ones((2, 2))\n",
        "linear_layer = Linear(4, 2)\n",
        "y = linear_layer(x)\n",
        "print(y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "382960020a56"
      },
      "source": [
        "Note that the weights `w` and `b` are automatically tracked by the layer upon\n",
        "being set as layer attributes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d3d875af9465"
      },
      "outputs": [],
      "source": [
        "assert linear_layer.weights == [linear_layer.w, linear_layer.b]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ec9d72aa7538"
      },
      "source": [
        "Note you also have access to a quicker shortcut for adding weight to a layer:\n",
        "the `add_weight()` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "168548eba841"
      },
      "outputs": [],
      "source": [
        "class Linear(keras.layers.Layer):\n",
        "    def __init__(self, units=32, input_dim=32):\n",
        "        super(Linear, self).__init__()\n",
        "        self.w = self.add_weight(\n",
        "            shape=(input_dim, units), initializer=\"random_normal\", trainable=True\n",
        "        )\n",
        "        self.b = self.add_weight(shape=(units,), initializer=\"zeros\", trainable=True)\n",
        "\n",
        "    def call(self, inputs):\n",
        "        return tf.matmul(inputs, self.w) + self.b\n",
        "\n",
        "\n",
        "x = tf.ones((2, 2))\n",
        "linear_layer = Linear(4, 2)\n",
        "y = linear_layer(x)\n",
        "print(y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "070ea9b4db6c"
      },
      "source": [
        "## Layers can have non-trainable weights\n",
        "\n",
        "Besides trainable weights, you can add non-trainable weights to a layer as\n",
        "well. Such weights are meant not to be taken into account during\n",
        "backpropagation, when you are training the layer.\n",
        "\n",
        "Here's how to add and use a non-trainable weight:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7c4cb404145f"
      },
      "outputs": [],
      "source": [
        "class ComputeSum(keras.layers.Layer):\n",
        "    def __init__(self, input_dim):\n",
        "        super(ComputeSum, self).__init__()\n",
        "        self.total = tf.Variable(initial_value=tf.zeros((input_dim,)), trainable=False)\n",
        "\n",
        "    def call(self, inputs):\n",
        "        self.total.assign_add(tf.reduce_sum(inputs, axis=0))\n",
        "        return self.total\n",
        "\n",
        "\n",
        "x = tf.ones((2, 2))\n",
        "my_sum = ComputeSum(2)\n",
        "y = my_sum(x)\n",
        "print(y.numpy())\n",
        "y = my_sum(x)\n",
        "print(y.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40f5b74d3d87"
      },
      "source": [
        "It's part of `layer.weights`, but it gets categorized as a non-trainable weight:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3d4db4ef4fa4"
      },
      "outputs": [],
      "source": [
        "print(\"weights:\", len(my_sum.weights))\n",
        "print(\"non-trainable weights:\", len(my_sum.non_trainable_weights))\n",
        "\n",
        "# It's not included in the trainable weights:\n",
        "print(\"trainable_weights:\", my_sum.trainable_weights)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fe6942aff7c6"
      },
      "source": [
        "## Best practice: deferring weight creation until the shape of the inputs is known\n",
        "\n",
        "Our `Linear` layer above took an `input_dim `argument that was used to compute\n",
        "the shape of the weights `w` and `b` in `__init__()`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "275b68d5ea9f"
      },
      "outputs": [],
      "source": [
        "class Linear(keras.layers.Layer):\n",
        "    def __init__(self, units=32, input_dim=32):\n",
        "        super(Linear, self).__init__()\n",
        "        self.w = self.add_weight(\n",
        "            shape=(input_dim, units), initializer=\"random_normal\", trainable=True\n",
        "        )\n",
        "        self.b = self.add_weight(shape=(units,), initializer=\"zeros\", trainable=True)\n",
        "\n",
        "    def call(self, inputs):\n",
        "        return tf.matmul(inputs, self.w) + self.b\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ebcacebb348"
      },
      "source": [
        "In many cases, you may not know in advance the size of your inputs, and you\n",
        "would like to lazily create weights when that value becomes known, some time\n",
        "after instantiating the layer.\n",
        "\n",
        "In the Keras API, we recommend creating layer weights in the `build(self,\n",
        "inputs_shape)` method of your layer. Like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "118c899f427e"
      },
      "outputs": [],
      "source": [
        "class Linear(keras.layers.Layer):\n",
        "    def __init__(self, units=32):\n",
        "        super(Linear, self).__init__()\n",
        "        self.units = units\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        self.w = self.add_weight(\n",
        "            shape=(input_shape[-1], self.units),\n",
        "            initializer=\"random_normal\",\n",
        "            trainable=True,\n",
        "        )\n",
        "        self.b = self.add_weight(\n",
        "            shape=(self.units,), initializer=\"random_normal\", trainable=True\n",
        "        )\n",
        "\n",
        "    def call(self, inputs):\n",
        "        return tf.matmul(inputs, self.w) + self.b\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "78061e0583c6"
      },
      "source": [
        "The `__call__()` method of your layer will automatically run build the first time\n",
        "it is called. You now have a layer that's lazy and thus easier to use:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fb08b1a45d22"
      },
      "outputs": [],
      "source": [
        "# At instantiation, we don't know on what inputs this is going to get called\n",
        "linear_layer = Linear(32)\n",
        "\n",
        "# The layer's weights are created dynamically the first time the layer is called\n",
        "y = linear_layer(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ddd7e8b22641"
      },
      "source": [
        "## Layers are recursively composable\n",
        "\n",
        "If you assign a Layer instance as an attribute of another Layer, the outer layer\n",
        "will start tracking the weights of the inner layer.\n",
        "\n",
        "We recommend creating such sublayers in the `__init__()` method (since the\n",
        "sublayers will typically have a build method, they will be built when the\n",
        "outer layer gets built)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9561cbf2fc60"
      },
      "outputs": [],
      "source": [
        "# Let's assume we are reusing the Linear class\n",
        "# with a `build` method that we defined above.\n",
        "\n",
        "\n",
        "class MLPBlock(keras.layers.Layer):\n",
        "    def __init__(self):\n",
        "        super(MLPBlock, self).__init__()\n",
        "        self.linear_1 = Linear(32)\n",
        "        self.linear_2 = Linear(32)\n",
        "        self.linear_3 = Linear(1)\n",
        "\n",
        "    def call(self, inputs):\n",
        "        x = self.linear_1(inputs)\n",
        "        x = tf.nn.relu(x)\n",
        "        x = self.linear_2(x)\n",
        "        x = tf.nn.relu(x)\n",
        "        return self.linear_3(x)\n",
        "\n",
        "\n",
        "mlp = MLPBlock()\n",
        "y = mlp(tf.ones(shape=(3, 64)))  # The first call to the `mlp` will create the weights\n",
        "print(\"weights:\", len(mlp.weights))\n",
        "print(\"trainable weights:\", len(mlp.trainable_weights))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "496736d98a62"
      },
      "source": [
        "## The `add_loss()` method\n",
        "\n",
        "When writing the `call()` method of a layer, you can create loss tensors that\n",
        "you will want to use later, when writing your training loop. This is doable by\n",
        "calling `self.add_loss(value)`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "084d56602ca4"
      },
      "outputs": [],
      "source": [
        "# A layer that creates an activity regularization loss\n",
        "class ActivityRegularizationLayer(keras.layers.Layer):\n",
        "    def __init__(self, rate=1e-2):\n",
        "        super(ActivityRegularizationLayer, self).__init__()\n",
        "        self.rate = rate\n",
        "\n",
        "    def call(self, inputs):\n",
        "        self.add_loss(self.rate * tf.reduce_sum(inputs))\n",
        "        return inputs\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5009ff0d1feb"
      },
      "source": [
        "These losses (including those created by any inner layer) can be retrieved via\n",
        "`layer.losses`. This property is reset at the start of every `__call__()` to\n",
        "the top-level layer, so that `layer.losses` always contains the loss values\n",
        "created during the last forward pass."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5bff6d59aea7"
      },
      "outputs": [],
      "source": [
        "class OuterLayer(keras.layers.Layer):\n",
        "    def __init__(self):\n",
        "        super(OuterLayer, self).__init__()\n",
        "        self.activity_reg = ActivityRegularizationLayer(1e-2)\n",
        "\n",
        "    def call(self, inputs):\n",
        "        return self.activity_reg(inputs)\n",
        "\n",
        "\n",
        "layer = OuterLayer()\n",
        "assert len(layer.losses) == 0  # No losses yet since the layer has never been called\n",
        "\n",
        "_ = layer(tf.zeros(1, 1))\n",
        "assert len(layer.losses) == 1  # We created one loss value\n",
        "\n",
        "# `layer.losses` gets reset at the start of each __call__\n",
        "_ = layer(tf.zeros(1, 1))\n",
        "assert len(layer.losses) == 1  # This is the loss created during the call above"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "36751ebe3363"
      },
      "source": [
        "In addition, the `loss` property also contains regularization losses created\n",
        "for the weights of any inner layer:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9327d3b581f8"
      },
      "outputs": [],
      "source": [
        "class OuterLayerWithKernelRegularizer(keras.layers.Layer):\n",
        "    def __init__(self):\n",
        "        super(OuterLayerWithKernelRegularizer, self).__init__()\n",
        "        self.dense = keras.layers.Dense(\n",
        "            32, kernel_regularizer=tf.keras.regularizers.l2(1e-3)\n",
        "        )\n",
        "\n",
        "    def call(self, inputs):\n",
        "        return self.dense(inputs)\n",
        "\n",
        "\n",
        "layer = OuterLayerWithKernelRegularizer()\n",
        "_ = layer(tf.zeros((1, 1)))\n",
        "\n",
        "# This is `1e-3 * sum(layer.dense.kernel ** 2)`,\n",
        "# created by the `kernel_regularizer` above.\n",
        "print(layer.losses)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "99d502b8899c"
      },
      "source": [
        "These losses are meant to be taken into account when writing training loops,\n",
        "like this:\n",
        "\n",
        "```python\n",
        "# Instantiate an optimizer.\n",
        "optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)\n",
        "loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "\n",
        "# Iterate over the batches of a dataset.\n",
        "for x_batch_train, y_batch_train in train_dataset:\n",
        "  with tf.GradientTape() as tape:\n",
        "    logits = layer(x_batch_train)  # Logits for this minibatch\n",
        "    # Loss value for this minibatch\n",
        "    loss_value = loss_fn(y_batch_train, logits)\n",
        "    # Add extra losses created during this forward pass:\n",
        "    loss_value += sum(model.losses)\n",
        "\n",
        "  grads = tape.gradient(loss_value, model.trainable_weights)\n",
        "  optimizer.apply_gradients(zip(grads, model.trainable_weights))\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7fa2db4a631d"
      },
      "source": [
        "For a detailed guide about writing training loops, see the\n",
        "[guide to writing a training loop from scratch](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch/).\n",
        "\n",
        "These losses also work seamlessly with `fit()` (they get automatically summed\n",
        "and added to the main loss, if any):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7c534372bf8a"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "inputs = keras.Input(shape=(3,))\n",
        "outputs = ActivityRegularizationLayer()(inputs)\n",
        "model = keras.Model(inputs, outputs)\n",
        "\n",
        "# If there is a loss passed in `compile`, thee regularization\n",
        "# losses get added to it\n",
        "model.compile(optimizer=\"adam\", loss=\"mse\")\n",
        "model.fit(np.random.random((2, 3)), np.random.random((2, 3)))\n",
        "\n",
        "# It's also possible not to pass any loss in `compile`,\n",
        "# since the model already has a loss to minimize, via the `add_loss`\n",
        "# call during the forward pass!\n",
        "model.compile(optimizer=\"adam\")\n",
        "model.fit(np.random.random((2, 3)), np.random.random((2, 3)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c753fcbc1818"
      },
      "source": [
        "## The `add_metric()` method\n",
        "\n",
        "Similarly to `add_loss()`, layers also have an `add_metric()` method\n",
        "for tracking the moving average of a quantity during training.\n",
        "\n",
        "Consider the following layer: a \"logistic endpoint\" layer.\n",
        "It takes as inputs predictions & targets, it computes a loss which it tracks\n",
        "via `add_loss()`, and it computes an accuracy scalar, which it tracks via\n",
        "`add_metric()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "85dad61dc160"
      },
      "outputs": [],
      "source": [
        "class LogisticEndpoint(keras.layers.Layer):\n",
        "    def __init__(self, name=None):\n",
        "        super(LogisticEndpoint, self).__init__(name=name)\n",
        "        self.loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)\n",
        "        self.accuracy_fn = keras.metrics.BinaryAccuracy()\n",
        "\n",
        "    def call(self, targets, logits, sample_weights=None):\n",
        "        # Compute the training-time loss value and add it\n",
        "        # to the layer using `self.add_loss()`.\n",
        "        loss = self.loss_fn(targets, logits, sample_weights)\n",
        "        self.add_loss(loss)\n",
        "\n",
        "        # Log accuracy as a metric and add it\n",
        "        # to the layer using `self.add_metric()`.\n",
        "        acc = self.accuracy_fn(targets, logits, sample_weights)\n",
        "        self.add_metric(acc, name=\"accuracy\")\n",
        "\n",
        "        # Return the inference-time prediction tensor (for `.predict()`).\n",
        "        return tf.nn.softmax(logits)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cd8807cb9cbc"
      },
      "source": [
        "Metrics tracked in this way are accessible via `layer.metrics`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "475df7270265"
      },
      "outputs": [],
      "source": [
        "layer = LogisticEndpoint()\n",
        "\n",
        "targets = tf.ones((2, 2))\n",
        "logits = tf.ones((2, 2))\n",
        "y = layer(targets, logits)\n",
        "\n",
        "print(\"layer.metrics:\", layer.metrics)\n",
        "print(\"current accuracy value:\", float(layer.metrics[0].result()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9eda5113fd18"
      },
      "source": [
        "Just like for `add_loss()`, these metrics are tracked by `fit()`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "97f767613953"
      },
      "outputs": [],
      "source": [
        "inputs = keras.Input(shape=(3,), name=\"inputs\")\n",
        "targets = keras.Input(shape=(10,), name=\"targets\")\n",
        "logits = keras.layers.Dense(10)(inputs)\n",
        "predictions = LogisticEndpoint(name=\"predictions\")(logits, targets)\n",
        "\n",
        "model = keras.Model(inputs=[inputs, targets], outputs=predictions)\n",
        "model.compile(optimizer=\"adam\")\n",
        "\n",
        "data = {\n",
        "    \"inputs\": np.random.random((3, 3)),\n",
        "    \"targets\": np.random.random((3, 10)),\n",
        "}\n",
        "model.fit(data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4bdbac3f6c85"
      },
      "source": [
        "## You can optionally enable serialization on your layers\n",
        "\n",
        "If you need your custom layers to be serializable as part of a\n",
        "[Functional model](https://www.tensorflow.org/guide/keras/functional/), you can optionally implement a `get_config()`\n",
        "method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b359ed5289a8"
      },
      "outputs": [],
      "source": [
        "class Linear(keras.layers.Layer):\n",
        "    def __init__(self, units=32):\n",
        "        super(Linear, self).__init__()\n",
        "        self.units = units\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        self.w = self.add_weight(\n",
        "            shape=(input_shape[-1], self.units),\n",
        "            initializer=\"random_normal\",\n",
        "            trainable=True,\n",
        "        )\n",
        "        self.b = self.add_weight(\n",
        "            shape=(self.units,), initializer=\"random_normal\", trainable=True\n",
        "        )\n",
        "\n",
        "    def call(self, inputs):\n",
        "        return tf.matmul(inputs, self.w) + self.b\n",
        "\n",
        "    def get_config(self):\n",
        "        return {\"units\": self.units}\n",
        "\n",
        "\n",
        "# Now you can recreate the layer from its config:\n",
        "layer = Linear(64)\n",
        "config = layer.get_config()\n",
        "print(config)\n",
        "new_layer = Linear.from_config(config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "78b207f7acbc"
      },
      "source": [
        "Note that the `__init__()` method of the base `Layer` class takes some keyword\n",
        "arguments, in particular a `name` and a `dtype`. It's good practice to pass\n",
        "these arguments to the parent class in `__init__()` and to include them in the\n",
        "layer config:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "00a3432cd28c"
      },
      "outputs": [],
      "source": [
        "class Linear(keras.layers.Layer):\n",
        "    def __init__(self, units=32, **kwargs):\n",
        "        super(Linear, self).__init__(**kwargs)\n",
        "        self.units = units\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        self.w = self.add_weight(\n",
        "            shape=(input_shape[-1], self.units),\n",
        "            initializer=\"random_normal\",\n",
        "            trainable=True,\n",
        "        )\n",
        "        self.b = self.add_weight(\n",
        "            shape=(self.units,), initializer=\"random_normal\", trainable=True\n",
        "        )\n",
        "\n",
        "    def call(self, inputs):\n",
        "        return tf.matmul(inputs, self.w) + self.b\n",
        "\n",
        "    def get_config(self):\n",
        "        config = super(Linear, self).get_config()\n",
        "        config.update({\"units\": self.units})\n",
        "        return config\n",
        "\n",
        "\n",
        "layer = Linear(64)\n",
        "config = layer.get_config()\n",
        "print(config)\n",
        "new_layer = Linear.from_config(config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ad5d100cc969"
      },
      "source": [
        "If you need more flexibility when deserializing the layer from its config, you\n",
        "can also override the `from_config()` class method. This is the base\n",
        "implementation of `from_config()`:\n",
        "\n",
        "```python\n",
        "def from_config(cls, config):\n",
        "  return cls(**config)\n",
        "```\n",
        "\n",
        "To learn more about serialization and saving, see the complete\n",
        "[guide to saving and serializing models](https://www.tensorflow.org/guide/keras/save_and_serialize/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "741c6d134d65"
      },
      "source": [
        "## Privileged `training` argument in the `call()` method\n",
        "\n",
        "Some layers, in particular the `BatchNormalization` layer and the `Dropout`\n",
        "layer, have different behaviors during training and inference. For such\n",
        "layers, it is standard practice to expose a `training` (boolean) argument in\n",
        "the `call()` method.\n",
        "\n",
        "By exposing this argument in `call()`, you enable the built-in training and\n",
        "evaluation loops (e.g. `fit()`) to correctly use the layer in training and\n",
        "inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "67ca741d0cfb"
      },
      "outputs": [],
      "source": [
        "class CustomDropout(keras.layers.Layer):\n",
        "    def __init__(self, rate, **kwargs):\n",
        "        super(CustomDropout, self).__init__(**kwargs)\n",
        "        self.rate = rate\n",
        "\n",
        "    def call(self, inputs, training=None):\n",
        "        if training:\n",
        "            return tf.nn.dropout(inputs, rate=self.rate)\n",
        "        return inputs\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5284e22677da"
      },
      "source": [
        "## Privileged `mask` argument in the `call()` method\n",
        "\n",
        "The other privileged argument supported by `call()` is the `mask` argument.\n",
        "\n",
        "You will find it in all Keras RNN layers. A mask is a boolean tensor (one\n",
        "boolean value per timestep in the input) used to skip certain input timesteps\n",
        "when processing timeseries data.\n",
        "\n",
        "Keras will automatically pass the correct `mask` argument to `__call__()` for\n",
        "layers that support it, when a mask is generated by a prior layer.\n",
        "Mask-generating layers are the `Embedding`\n",
        "layer configured with `mask_zero=True`, and the `Masking` layer.\n",
        "\n",
        "To learn more about masking and how to write masking-enabled layers, please\n",
        "check out the guide\n",
        "[\"understanding padding and masking\"](https://www.tensorflow.org/guide/keras/masking_and_padding/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bf87358118de"
      },
      "source": [
        "## The `Model` class\n",
        "\n",
        "In general, you will use the `Layer` class to define inner computation blocks,\n",
        "and will use the `Model` class to define the outer model -- the object you\n",
        "will train.\n",
        "\n",
        "For instance, in a ResNet50 model, you would have several ResNet blocks\n",
        "subclassing `Layer`, and a single `Model` encompassing the entire ResNet50\n",
        "network.\n",
        "\n",
        "The `Model` class has the same API as `Layer`, with the following differences:\n",
        "\n",
        "- It exposes built-in training, evaluation, and prediction loops\n",
        "(`model.fit()`, `model.evaluate()`, `model.predict()`).\n",
        "- It exposes the list of its inner layers, via the `model.layers` property.\n",
        "- It exposes saving and serialization APIs (`save()`, `save_weights()`...)\n",
        "\n",
        "Effectively, the `Layer` class corresponds to what we refer to in the\n",
        "literature as a \"layer\" (as in \"convolution layer\" or \"recurrent layer\") or as\n",
        "a \"block\" (as in \"ResNet block\" or \"Inception block\").\n",
        "\n",
        "Meanwhile, the `Model` class corresponds to what is referred to in the\n",
        "literature as a \"model\" (as in \"deep learning model\") or as a \"network\" (as in\n",
        "\"deep neural network\").\n",
        "\n",
        "So if you're wondering, \"should I use the `Layer` class or the `Model` class?\",\n",
        "ask yourself: will I need to call `fit()` on it? Will I need to call `save()`\n",
        "on it? If so, go with `Model`. If not (either because your class is just a block\n",
        "in a bigger system, or because you are writing training & saving code yourself),\n",
        "use `Layer`.\n",
        "\n",
        "For instance, we could take our mini-resnet example above, and use it to build\n",
        "a `Model` that we could train with `fit()`, and that we could save with\n",
        "`save_weights()`:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "78539e55de0c"
      },
      "source": [
        "```python\n",
        "class ResNet(tf.keras.Model):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(ResNet, self).__init__()\n",
        "        self.block_1 = ResNetBlock()\n",
        "        self.block_2 = ResNetBlock()\n",
        "        self.global_pool = layers.GlobalAveragePooling2D()\n",
        "        self.classifier = Dense(num_classes)\n",
        "\n",
        "    def call(self, inputs):\n",
        "        x = self.block_1(inputs)\n",
        "        x = self.block_2(x)\n",
        "        x = self.global_pool(x)\n",
        "        return self.classifier(x)\n",
        "\n",
        "\n",
        "resnet = ResNet()\n",
        "dataset = ...\n",
        "resnet.fit(dataset, epochs=10)\n",
        "resnet.save(filepath)\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b817a4de8c5d"
      },
      "source": [
        "## Putting it all together: an end-to-end example\n",
        "\n",
        "Here's what you've learned so far:\n",
        "\n",
        "- A `Layer` encapsulate a state (created in `__init__()` or `build()`) and some\n",
        "computation (defined in `call()`).\n",
        "- Layers can be recursively nested to create new, bigger computation blocks.\n",
        "- Layers can create and track losses (typically regularization losses) as well\n",
        "as metrics, via `add_loss()` and `add_metric()`\n",
        "- The outer container, the thing you want to train, is a `Model`. A `Model` is\n",
        "just like a `Layer`, but with added training and serialization utilities.\n",
        "\n",
        "Let's put all of these things together into an end-to-end example: we're going\n",
        "to implement a Variational AutoEncoder (VAE). We'll train it on MNIST digits.\n",
        "\n",
        "Our VAE will be a subclass of `Model`, built as a nested composition of layers\n",
        "that subclass `Layer`. It will feature a regularization loss (KL divergence)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "18842173f875"
      },
      "outputs": [],
      "source": [
        "from tensorflow.keras import layers\n",
        "\n",
        "\n",
        "class Sampling(layers.Layer):\n",
        "    \"\"\"Uses (z_mean, z_log_var) to sample z, the vector encoding a digit.\"\"\"\n",
        "\n",
        "    def call(self, inputs):\n",
        "        z_mean, z_log_var = inputs\n",
        "        batch = tf.shape(z_mean)[0]\n",
        "        dim = tf.shape(z_mean)[1]\n",
        "        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))\n",
        "        return z_mean + tf.exp(0.5 * z_log_var) * epsilon\n",
        "\n",
        "\n",
        "class Encoder(layers.Layer):\n",
        "    \"\"\"Maps MNIST digits to a triplet (z_mean, z_log_var, z).\"\"\"\n",
        "\n",
        "    def __init__(self, latent_dim=32, intermediate_dim=64, name=\"encoder\", **kwargs):\n",
        "        super(Encoder, self).__init__(name=name, **kwargs)\n",
        "        self.dense_proj = layers.Dense(intermediate_dim, activation=\"relu\")\n",
        "        self.dense_mean = layers.Dense(latent_dim)\n",
        "        self.dense_log_var = layers.Dense(latent_dim)\n",
        "        self.sampling = Sampling()\n",
        "\n",
        "    def call(self, inputs):\n",
        "        x = self.dense_proj(inputs)\n",
        "        z_mean = self.dense_mean(x)\n",
        "        z_log_var = self.dense_log_var(x)\n",
        "        z = self.sampling((z_mean, z_log_var))\n",
        "        return z_mean, z_log_var, z\n",
        "\n",
        "\n",
        "class Decoder(layers.Layer):\n",
        "    \"\"\"Converts z, the encoded digit vector, back into a readable digit.\"\"\"\n",
        "\n",
        "    def __init__(self, original_dim, intermediate_dim=64, name=\"decoder\", **kwargs):\n",
        "        super(Decoder, self).__init__(name=name, **kwargs)\n",
        "        self.dense_proj = layers.Dense(intermediate_dim, activation=\"relu\")\n",
        "        self.dense_output = layers.Dense(original_dim, activation=\"sigmoid\")\n",
        "\n",
        "    def call(self, inputs):\n",
        "        x = self.dense_proj(inputs)\n",
        "        return self.dense_output(x)\n",
        "\n",
        "\n",
        "class VariationalAutoEncoder(keras.Model):\n",
        "    \"\"\"Combines the encoder and decoder into an end-to-end model for training.\"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        original_dim,\n",
        "        intermediate_dim=64,\n",
        "        latent_dim=32,\n",
        "        name=\"autoencoder\",\n",
        "        **kwargs\n",
        "    ):\n",
        "        super(VariationalAutoEncoder, self).__init__(name=name, **kwargs)\n",
        "        self.original_dim = original_dim\n",
        "        self.encoder = Encoder(latent_dim=latent_dim, intermediate_dim=intermediate_dim)\n",
        "        self.decoder = Decoder(original_dim, intermediate_dim=intermediate_dim)\n",
        "\n",
        "    def call(self, inputs):\n",
        "        z_mean, z_log_var, z = self.encoder(inputs)\n",
        "        reconstructed = self.decoder(z)\n",
        "        # Add KL divergence regularization loss.\n",
        "        kl_loss = -0.5 * tf.reduce_mean(\n",
        "            z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1\n",
        "        )\n",
        "        self.add_loss(kl_loss)\n",
        "        return reconstructed\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40384d934b3c"
      },
      "source": [
        "Let's write a simple training loop on MNIST:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c37fef01d4bc"
      },
      "outputs": [],
      "source": [
        "original_dim = 784\n",
        "vae = VariationalAutoEncoder(original_dim, 64, 32)\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
        "mse_loss_fn = tf.keras.losses.MeanSquaredError()\n",
        "\n",
        "loss_metric = tf.keras.metrics.Mean()\n",
        "\n",
        "(x_train, _), _ = tf.keras.datasets.mnist.load_data()\n",
        "x_train = x_train.reshape(60000, 784).astype(\"float32\") / 255\n",
        "\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices(x_train)\n",
        "train_dataset = train_dataset.shuffle(buffer_size=1024).batch(64)\n",
        "\n",
        "epochs = 2\n",
        "\n",
        "# Iterate over epochs.\n",
        "for epoch in range(epochs):\n",
        "    print(\"Start of epoch %d\" % (epoch,))\n",
        "\n",
        "    # Iterate over the batches of the dataset.\n",
        "    for step, x_batch_train in enumerate(train_dataset):\n",
        "        with tf.GradientTape() as tape:\n",
        "            reconstructed = vae(x_batch_train)\n",
        "            # Compute reconstruction loss\n",
        "            loss = mse_loss_fn(x_batch_train, reconstructed)\n",
        "            loss += sum(vae.losses)  # Add KLD regularization loss\n",
        "\n",
        "        grads = tape.gradient(loss, vae.trainable_weights)\n",
        "        optimizer.apply_gradients(zip(grads, vae.trainable_weights))\n",
        "\n",
        "        loss_metric(loss)\n",
        "\n",
        "        if step % 100 == 0:\n",
        "            print(\"step %d: mean loss = %.4f\" % (step, loss_metric.result()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "65e5faeb0029"
      },
      "source": [
        "Note that since the VAE is subclassing `Model`, it features built-in training\n",
        "loops. So you could also have trained it like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1e98ba7ebdb8"
      },
      "outputs": [],
      "source": [
        "vae = VariationalAutoEncoder(784, 64, 32)\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
        "\n",
        "vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())\n",
        "vae.fit(x_train, x_train, epochs=2, batch_size=64)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "42ee3169e70c"
      },
      "source": [
        "## Beyond object-oriented development: the Functional API\n",
        "\n",
        "Was this example too much object-oriented development for you? You can also\n",
        "build models using the [Functional API](https://www.tensorflow.org/guide/keras/functional/). Importantly,\n",
        "choosing one style or another does not prevent you from leveraging components\n",
        "written in the other style: you can always mix-and-match.\n",
        "\n",
        "For instance, the Functional API example below reuses the same `Sampling` layer\n",
        "we defined in the example above:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b8fe39f892c7"
      },
      "outputs": [],
      "source": [
        "original_dim = 784\n",
        "intermediate_dim = 64\n",
        "latent_dim = 32\n",
        "\n",
        "# Define encoder model.\n",
        "original_inputs = tf.keras.Input(shape=(original_dim,), name=\"encoder_input\")\n",
        "x = layers.Dense(intermediate_dim, activation=\"relu\")(original_inputs)\n",
        "z_mean = layers.Dense(latent_dim, name=\"z_mean\")(x)\n",
        "z_log_var = layers.Dense(latent_dim, name=\"z_log_var\")(x)\n",
        "z = Sampling()((z_mean, z_log_var))\n",
        "encoder = tf.keras.Model(inputs=original_inputs, outputs=z, name=\"encoder\")\n",
        "\n",
        "# Define decoder model.\n",
        "latent_inputs = tf.keras.Input(shape=(latent_dim,), name=\"z_sampling\")\n",
        "x = layers.Dense(intermediate_dim, activation=\"relu\")(latent_inputs)\n",
        "outputs = layers.Dense(original_dim, activation=\"sigmoid\")(x)\n",
        "decoder = tf.keras.Model(inputs=latent_inputs, outputs=outputs, name=\"decoder\")\n",
        "\n",
        "# Define VAE model.\n",
        "outputs = decoder(z)\n",
        "vae = tf.keras.Model(inputs=original_inputs, outputs=outputs, name=\"vae\")\n",
        "\n",
        "# Add KL divergence regularization loss.\n",
        "kl_loss = -0.5 * tf.reduce_mean(z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)\n",
        "vae.add_loss(kl_loss)\n",
        "\n",
        "# Train.\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\n",
        "vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())\n",
        "vae.fit(x_train, x_train, epochs=3, batch_size=64)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0f4db7df7eb5"
      },
      "source": [
        "For more information, make sure to read the [Functional API guide](https://www.tensorflow.org/guide/keras/functional/)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "custom_layers_and_models.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
